import {CodeClient} from './CodeClient';
import {CodeFold} from './CodeFold';
import {CodeLink} from './Link';
import {CodeProps} from './VisualExampleClient';
import {HastNode, HastTextNode, highlightHast, Language} from 'tree-sitter-highlight';
import React, {cache} from 'react';
import {style, StyleString} from '@react-spectrum/s2/style' with {type: 'macro'};
import {TabLink} from './FileTabs';
import {Token, TokenType} from './CodeToken';

const property = style({color: 'indigo-1000'});
const fn = style({color: 'red-1000'});
export const styles = {
  keyword: style({color: 'magenta-1000'}),
  string: style({color: 'green-1000'}),
  number: style({color: 'pink-1000'}),
  property,
  attribute: property,
  function: fn,
  tag: fn,
  constructor: fn,
  comment: style({color: 'gray-700'}),
  variable: style({color: 'fuchsia-1000'}),
  import: style({display: '--import-display'})
};

const mark = style({
  display: 'block',
  backgroundColor: 'blue-800/15',
  borderColor: 'blue-800',
  borderWidth: 0,
  borderStartWidth: 2,
  borderStyle: 'solid',
  marginStart: 'calc(var(--code-padding-start) * -1)',
  marginEnd: 'calc(var(--code-padding-end) * -1)',
  paddingStart: 'calc(var(--code-padding-start) - self(borderStartWidth))',
  paddingEnd: '--code-padding-end',
  color: 'inherit'
});

function Highlight({tokens}) {
  return <mark className={mark}><CodeClient tokens={tokens} /></mark>;
}

function Focus({tokens}) {
  return <span><CodeClient tokens={tokens} /></span>;
}

const groupings = {
  highlight: Highlight,
  collapse: CodeFold,
  focus: Focus
};

type Links = {[name: string]: string};
export interface ICodeProps {
  children: string,
  lang?: string,
  hideImports?: boolean,
  links?: Links,
  styles?: StyleString
}

export function Code({children, lang, hideImports = true, links, styles}: ICodeProps) {
  if (lang) {
    return (
      <code className={styles} style={{fontFamily: 'inherit', WebkitTextSizeAdjust: 'none'}}>
        <CodeClient tokens={highlightCode(children, lang, hideImports, links)} />
      </code>
    );
  }

  return (
    <code
      className={style({
        font: 'code-sm',
        backgroundColor: 'layer-1',
        paddingX: 4,
        borderWidth: 1,
        borderColor: 'gray-100',
        borderStyle: 'solid',
        borderRadius: 'sm',
        whiteSpace: 'pre-wrap'
      })}>
      {children}
    </code>
  );
}

const highlightCode = cache((children: string, lang: string, hideImports = true, links?: Links): Token[] => {
  // @ts-ignore
  let highlighted = highlightHast(children, Language[lang === 'json' ? 'JS' : lang.toUpperCase()]);
  let lineNodes = lines(highlighted);
  let idx = lineNodes.findIndex(line => !/^(["']use client["']|(\s*$))/.test(text(line)));
  if (idx > 0) {
    lineNodes = lineNodes.slice(idx);
  }

  if (hideImports) {
    // Group into hidden and visible nodes.
    // Hidden nodes will include all import statements. If a highlighted block is seen,
    // then we'll hide all the lines up until 2 lines before this.
    let hidden: HastNode[] = [];
    let visible: HastNode[] = [];
    let seenNonImportLine = false;
    let hasHighlight = false;
    for (let line of lineNodes) {
      if (!seenNonImportLine && /^(["']use client["']|@?import|(\s*$))/.test(text(line))) {
        hidden.push(line);
      } else {
        seenNonImportLine = true;
        visible.push(line);
      }

      if ((line.tagName === 'highlight' || line.tagName === 'focus') && !hasHighlight) {
        hasHighlight = true;
        // Center highlighted lines within collapsed window (~8 lines).
        let highlightedLines = line.children.length;
        let contextLines = highlightedLines < 6
          ? Math.floor((8 - highlightedLines) / 2)
          : 2;
        contextLines++;
        hidden.push(...visible.slice(0, -contextLines));
        visible = visible.slice(-contextLines);
      }
    }

    if (hidden.length && visible.length) {
      lineNodes = [
        {
          type: 'element',
          tagName: 'span',
          children: hidden,
          properties: {
            className: 'import'
          }
        },
        ...visible
      ];
    }
  }

  return renderChildren(lineNodes, '0', links);
});

function lines(node: HastNode) {
  let resultLines: HastNode[] = [];
  let currentLine: (HastNode | HastTextNode)[] = [];
  // let properties: Record<string, string> = {};
  let grouping: HastNode | null = null;
  let skip = false;
  let endLine = () => {
    if (skip) {
      skip = false;
      currentLine = [];
      return;
    }
    let childNode = {
      type: 'element',
      tagName: 'div',
      children: [{...node, children: currentLine}]
    } as HastNode;
    if (grouping) {
      grouping.children.push(childNode);
    } else {
      resultLines.push(childNode);
    }
    currentLine = [];
  };

  for (let child of node.children) {
    if (child.type === 'text' && 'value' in child) {
      let parts = child.value.split('\n');
      for (let part of parts.slice(0, -1)) {
        if (part.length) {
          currentLine.push({type: 'text', value: part});
        }
        endLine();
      }
      let last = parts.at(-1);
      if (last?.length) {
        currentLine.push({type: 'text', value: last});
      }
      continue;
    } else if ('properties' in child && child.properties?.className === 'comment') {
      let comment = text(child);
      let begin = comment.match(/^(?:\/\/\/|\/\*)- begin (.+) -(?:\/\/\/|\*\/)$/);
      if (begin) {
        grouping = {
          type: 'element',
          tagName: begin[1],
          children: []
        } as any;
        currentLine = [];
        skip = true;
        continue;
      } else if (grouping && (comment.startsWith('///- end ') || comment.startsWith('/*- end '))) {
        resultLines.push(grouping);
        grouping = null;
        currentLine = [];
        skip = true;
        continue;
      }
    }

    let result = lines(child as HastNode);
    if (result.length) {
      currentLine.push(...result[0].children);
    }
    if (result.length > 1) {
      endLine();
      for (let i = 1; i < result.length - 1; i++) {
        currentLine = result[i].children;
        endLine();
      }
      currentLine = result.at(-1)!.children;
    }
  }

  if (currentLine.length) {
    endLine();
  }

  return resultLines;
}

// Renders a Hast Node to a list of tokens. A token is either a string, a React element, or a token type (number) + string.
// These are flattened into an array that gets sent to the client. This format significantly reduces the payload size vs JSX.
function renderHast(node: HastNode | HastTextNode, key: string, links?: Links, indent = ''): Token | Token[] {
  if (node.type === 'element' && 'children' in node) {
    let childArray: Token[] = renderChildren(node.children, key, links);
    if (node.tagName === 'div') {
      if (typeof childArray.at(-1) === 'string' && typeof childArray.at(-2) !== 'number') {
        childArray[childArray.length - 1] += '\n';
      } else {
        childArray.push('\n');
      }
    }

    let tokenType = node.properties?.className.split(' ').map(c => TokenType[c]).filter(v => v != null) || [];
    if (node.properties?.className === 'comment' && text(node) === '/* PROPS */') {
      return <CodeProps key={key} indent={indent} />;
    }

    // CodeProps includes the indent and newlines in case there are no props to show.
    if (node.tagName === 'div' && typeof childArray[0] === 'string' && /^\s+$/.test(childArray[0]) && React.isValidElement(childArray[1]) && childArray[1].type === CodeProps) {
      // If the only thing after CodeProps is the newline from div processing, exclude it (CodeProps handles its own newlines).
      // Otherwise, include all trailing content.
      childArray = childArray.length === 3 && childArray[2] === '\n'
      ? childArray.slice(1, 2)
      : childArray.slice(1);
    }

    let children = childArray.length === 1 ? childArray[0] : childArray;
    let tagName: any = node.tagName;
    let properties: any = node.properties;
    if (links && typeof children === 'string' && links[children]) {
      let link = links[children];
      return (
        <CodeLink
          key={key}
          className={styles[properties?.className]}
          href={link}>
          {children}
        </CodeLink>
      );
    }

    // Link to imported files.
    if (properties?.className === 'string' && typeof children === 'string' && /^['"]\.\//.test(children)) {
      return (
        <TabLink
          key={key}
          className={styles.string}
          name={children.slice(3, -1)}>
          <CodeClient tokens={childArray} />
        </TabLink>
      );
    }

    if ((tagName === 'div' || tagName === 'span') && tokenType.length === 0) {
      return children;
    }

    let type = groupings[tagName] || tagName;
    if (type === 'div') {
      // we inserted a newline, so treat this as inline.
      type = 'span';
    }

    if (type === 'span') {
      return [tokenType[0], children];
    }

    let className = node.properties?.className.split(' ').map(c => styles[c]).filter(Boolean).join(' ') || undefined;
    return React.createElement(type, {...properties, className, key, tokens: childArray});
  } else {
    // @ts-ignore
    return node.value;
  }
}

function renderChildren(children: (HastNode | HastTextNode)[], key: string, links?: Links): Token[] {
  let childArray: Token[] = [];
  let type = -1;
  let stringIndex = -1;
  for (let [i, child] of children.entries()) {
    let indent = i === 1 && stringIndex >= 0 && /^\s+$/.test(childArray[stringIndex] as string) ? childArray[stringIndex] as string : '';
    let childNode = renderHast(child, `${key}.${i}`, links, indent);
    let childNodes = Array.isArray(childNode) ? childNode : [childNode];
    let childIndex = 0;
    while (childIndex < childNodes.length) {
      let child = childNodes[childIndex++];
      let childType = -1;
      if (typeof child === 'number') {
        // A number represents a token type. Consume the next value.
        childType = child;
        child = childNodes[childIndex++];
        if (childType !== type) {
          childArray.push(childType);
        }
      }

      // If this is a string, either append to the previous string if it is
      // the same token type, or push a new string.
      if (typeof child === 'string') {
        if (childType !== type) {
          type = childType;
          stringIndex = childArray.length;
          childArray.push(child);
        } else if (stringIndex >= 0) {
          childArray[stringIndex] += child;
        } else {
          stringIndex = childArray.length;
          childArray.push(child);
        }
      } else {
        type = -1;
        stringIndex = -1;
        childArray.push(child);
      }
    }
  }

  return childArray;
}

function text(node) {
  if (node.type === 'text') {
    return node.value;
  } else {
    return node.children.map(c => text(c)).join('');
  }
}
